import os
import cPickle
import logging
import numpy as np
import pandas as pd
from PIL import Image as PILImage
# import Image
import cStringIO as StringIO

import sys

sys.path.append('/home/dell/wxm/Code/KaggleDDD/Caffe_CRFasRNN/caffe/python')
import caffe
import numpy as np
import os
from skimage import io
import matplotlib.pyplot as plt
import cv2

MODEL_FILE = 'TVG_CRFRNN_new_deploy.prototxt'
PRETRAINED = 'TVG_CRFRNN_COCO_VOC.caffemodel'
DEBUG = True
caffe.set_mode_gpu()
caffe.set_device(1)
net = caffe.Segmenter(MODEL_FILE, PRETRAINED, False)

img_folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/KaggleDDD/train'
img_folder_seg = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/KaggleDDD/seg_train'

mean_vec = np.array([103.939, 116.779, 123.68], dtype=np.float32)
reshaped_mean_vec = mean_vec.reshape(1, 1, 3)

for c in os.listdir(img_folder):
    flag = 0
    img_folder_c = img_folder + '/' + c
    img_folder_c_seg = img_folder_seg + '/seg_' + c
    if os.path.exists(img_folder_c_seg) is False:
        os.mkdir(img_folder_c_seg)
    for img in os.listdir(img_folder_c):
        img_path = img_folder_c + '/' + img
        img_path_seg = img_folder_c_seg + '/' + img
        input_im = io.imread(img_path)
        input_im = cv2.resize(input_im, (500, 375), interpolation=cv2.INTER_AREA)
        # Rearrange channels to form BGR
        im = input_im[:, :, ::-1]
        # Subtract mean
        im = im - reshaped_mean_vec
        # Pad as necessary
        cur_h, cur_w, cur_c = im.shape
        pad_h = 500 - cur_h
        pad_w = 500 - cur_w
        im = np.pad(im, pad_width=((0, pad_h), (0, pad_w), (0, 0)), mode='constant', constant_values=0)
        # Get predictions
        segmentation = net.predict([im])
        segmentation2 = segmentation[0:cur_h, 0:cur_w]
        segmentation2[segmentation2 == 15] = 255
        io.imsave(fname=img_path_seg, arr=segmentation2)

        flag += 1
        print c
        print flag
        if DEBUG is True and flag >= 10:
            break
